# 防止重新下载东西
import os
from io import StringIO
from io import BytesIO
from torch import tensor 
if os.name == 'nt':
    os.chdir("C:\\federated_malware\\Project")
    print ("cwd",os.getcwd())#获得当前目录
    print ("工作目录",os.path.abspath('.'))#获得当前工作目录
    print ("工作目录",os.path.abspath(os.curdir))#获得当前工作目录

import sys
sys.path.append('C:\\federated_malware\\Project\\')
sys.path.append('C:\\federated_malware\\Project\\rpc\\grpc_gen')
# print(sys.path)

import logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)


import grpc
from rpc.grpc_gen import trainer_pb2 as pb2
from rpc.grpc_gen import trainer_pb2_grpc as pb2_grpc

from trainers.fedbase.FedClientBase import FedClientBase
from models.minist.dnn_model import Net
import torch
# 需要将原来的请求，转换成grpc的请求。即输入参数和输出参数转换回来。内部使用远程调用。
# 用来生成虚拟的远程客户端。继承原来的客户端
class FedProxyClient(FedClientBase):
    def __init__(self,worker_id='localhost:50051'):
        # 初始化远程调用的一些东西。
        # worker_id默认是远程调用的worker的IP地址
        self.worker_id = worker_id
        channel = grpc.insecure_channel(worker_id) # 连接上gRPC服务端
        self.stub = pb2_grpc.trainerStub(channel)

    def set_parameters(self,client_params):
        # 调用远程方法.相当于给类添加一个属性
        request = pb2.set_parameters_request(
            epochs=client_params['epochs'],
            batch_size=client_params['batch_size'],
            inner_lr=client_params['inner_lr'],
            outer_lr=client_params['outer_lr'],
            update_steps=client_params['update_steps']
        )
        # 验证参数。命名参数没有办法直接传字典。
        # request = pb2.set_parameters_request(client_params)
        response = self.stub.set_parameters(request)
        # return response

    def process_data(self):
        request = pb2.empty_request()
        response = self.stub.process_data(request)
        # return response

    def process_data_support_and_query(self):
        request = pb2.empty_request()
        response = self.stup.process_data_support_and_query(request)
        # return response

    def load_model_state_dict(self,model_state_dict):
        # request = pb2.model_state_dict_request(model_state_dict)
        # 验证了映射参数，可以直接传递字典。
        bytes_buffer = BytesIO()
        torch.save(model_state_dict,bytes_buffer)
        request = pb2.model_state_dict_request(state_dict = bytes_buffer.getvalue())
        response = self.stub.load_model_state_dict(request)
        # return response


    def get_model_state_dict(self):
        request = pb2.empty_request()
        response = self.stub.get_model_state_dict(request)
        bytes_buffer = BytesIO(response.state_dict)
        model_state_dict = torch.load(bytes_buffer)
        return model_state_dict

    def save_model(self):
        request = pb2.empty_request()
        response = self.save_model(request)
        # return response

    def read_model(self):
        request = pb2.empty_request()
        response = self.stub.get_model(request)
        # return response

    def train(self):
        request = pb2.empty_request()
        response = self.stub.train(request)
        return response.result
    
    def test(self):
        request = pb2.empty_request()
        response = self.stub.test(request)
        return response.result



def worker_test():
    worker = FedProxyClient()
    client_params = {
            # 迭代轮次
            "epochs" :1,
            # batchsize 被update_steps代替了
            "batch_size":128,
            # 本地梯度下降的步长
            "inner_lr":0.01,
            # 服务器梯度下降的步长
            "outer_lr":0.8,
            # 本地更新的次数，会覆盖batch_size。batch_size=len(dataset)/update_steps
            "update_steps": 10,
        }
    # 测试远程代理——设置参数
    worker.set_parameters(client_params)
    # 测试远程代理——处理数据
    worker.process_data()
    # 测试远程代理——加载参数
    model = Net()
    worker.load_model_state_dict(model.state_dict())
    # 测试远程代理——获取参数
    worker.get_model_state_dict()
    #  测试远程代理——训练
    worker.train()
    # 测试远程代理——测试
    worker.test()


if __name__ == "__main__":
    worker_test()